"""
示例 3: 创建圆柱对称电极

演示如何：
1. 使用圆柱对称
2. 创建圆柱形电极
3. 创建孔径和管道结构
"""

import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from SIMION.PA import PA


def create_cylindrical_lens():
    """创建圆柱对称透镜"""
    
    print("\n" + "="*60)
    print("创建圆柱对称透镜")
    print("="*60 + "\n")
    
    # 创建圆柱对称阵列
    # X 轴表示径向距离（r）
    # Y 轴表示轴向距离（z）
    pa = PA(
        nx=50,                     # 径向 50 网格点
        ny=100,                    # 轴向 100 网格点
        nz=1,                      # 不使用 Z
        symmetry='cylindrical',    # 圆柱对称
        max_voltage=5000,          # 最大电压 5kV
        field_type='electrostatic',
        dx_mm=0.5,                 # 每个网格 0.5mm
        dy_mm=0.5
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} (圆柱对称)")
    print(f"网格尺寸: {pa.dx_mm()} mm x {pa.dy_mm()} mm")
    print(f"实际尺寸: {pa.nx()*pa.dx_mm()} mm (r) x {pa.ny()*pa.dy_mm()} mm (z)\n")
    
    # 创建管道结构
    outer_radius = 40  # 外半径（网格单位）
    inner_radius = 15  # 内半径（网格单位）
    
    print(f"管道参数:")
    print(f"  外半径: {outer_radius} 网格 ({outer_radius*pa.dx_mm()} mm)")
    print(f"  内半径: {inner_radius} 网格 ({inner_radius*pa.dx_mm()} mm)\n")
    
    # 入口电极 (0V)
    print("创建入口电极...")
    entrance_start = 10
    entrance_end = 15
    for y in range(entrance_start, entrance_end + 1):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 0.0)
    
    # 加速电极 (+1000V)
    print("创建加速电极...")
    accel_start = 40
    accel_end = 60
    for y in range(accel_start, accel_end + 1):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 1000.0)
    
    # 出口电极 (0V)
    print("创建出口电极...")
    exit_start = 85
    exit_end = 90
    for y in range(exit_start, exit_end + 1):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 0.0)
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "cylindrical_lens.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print(f"\n文件已保存: cylindrical_lens.pa#\n")
    
    return pa


def create_aperture():
    """创建孔径电极"""
    
    print("\n" + "="*60)
    print("创建孔径电极")
    print("="*60 + "\n")
    
    # 创建小型圆柱对称阵列
    pa = PA(
        nx=30,                     # 径向 30 网格点
        ny=30,                     # 轴向 30 网格点
        nz=1,
        symmetry='cylindrical',
        max_voltage=1000,
        field_type='electrostatic',
        dx_mm=0.2,
        dy_mm=0.2
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} (圆柱对称)")
    print(f"实际尺寸: {pa.nx()*pa.dx_mm()} mm x {pa.ny()*pa.dy_mm()} mm\n")
    
    # 创建薄孔径
    aperture_thickness = 2  # 孔径厚度（网格单位）
    aperture_radius = 5     # 孔径半径
    outer_radius = 25       # 外半径
    
    aperture_center_y = pa.ny() // 2
    aperture_y_start = aperture_center_y - aperture_thickness // 2
    aperture_y_end = aperture_center_y + aperture_thickness // 2
    
    print(f"孔径参数:")
    print(f"  孔径半径: {aperture_radius} 网格 ({aperture_radius*pa.dx_mm()} mm)")
    print(f"  孔径厚度: {aperture_thickness} 网格 ({aperture_thickness*pa.dy_mm()} mm)")
    print(f"  孔径位置: Y={aperture_center_y}\n")
    
    print("创建孔径电极 (100V)...")
    for y in range(aperture_y_start, aperture_y_end + 1):
        for r in range(aperture_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 100.0)
    
    # 验证孔径中心
    is_elec_inside, pot_inside = pa.point(aperture_radius - 1, aperture_center_y, 0)
    is_elec_on, pot_on_aperture = pa.point(aperture_radius + 1, aperture_center_y, 0)
    
    print(f"\n验证:")
    print(f"  孔径内部 (r={aperture_radius-1}): 电位={pot_inside}V, 电极={is_elec_inside}")
    print(f"  孔径上 (r={aperture_radius+1}): 电位={pot_on_aperture}V, 电极={is_elec_on}")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "aperture.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print(f"\n文件已保存: aperture.pa#\n")
    
    return pa


def create_ion_funnel():
    """创建离子漏斗结构"""
    
    print("\n" + "="*60)
    print("创建离子漏斗")
    print("="*60 + "\n")
    
    pa = PA(
        nx=60,
        ny=120,
        nz=1,
        symmetry='cylindrical',
        max_voltage=2000,
        field_type='electrostatic',
        dx_mm=0.5,
        dy_mm=0.5
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} (圆柱对称)\n")
    
    # 创建一系列环形电极，半径逐渐减小
    print("创建漏斗形环形电极...")
    
    num_rings = 10
    ring_thickness = 2
    ring_spacing = 10
    
    for i in range(num_rings):
        y_pos = 10 + i * ring_spacing
        
        # 半径从大到小
        ring_radius = 50 - i * 4
        ring_outer = ring_radius + 3
        
        # 电压逐渐增加
        voltage = i * 20.0
        
        print(f"  环 {i+1}: Y={y_pos}, R={ring_radius}, V={voltage}V")
        
        # 创建环形电极
        for y in range(y_pos, y_pos + ring_thickness):
            for r in range(ring_radius, min(ring_outer, pa.nx())):
                pa.point(r, y, 0, 1, voltage)
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "ion_funnel.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print(f"\n文件已保存: ion_funnel.pa#\n")
    
    return pa


if __name__ == "__main__":
    # 创建各种圆柱对称结构
    print("\n" + "="*60)
    print("SIMION 圆柱对称电极创建示例")
    print("="*60)
    
    # 1. 圆柱对称透镜
    pa1 = create_cylindrical_lens()
    
    # 2. 孔径
    pa2 = create_aperture()
    
    # 3. 离子漏斗
    pa3 = create_ion_funnel()
    
    print("\n" + "="*60)
    print("所有圆柱对称结构创建完成!")
    print("="*60)
    print("\n生成的文件:")
    print("  - cylindrical_lens.pa# : 圆柱对称透镜")
    print("  - aperture.pa# : 孔径电极")
    print("  - ion_funnel.pa# : 离子漏斗")
    print("\n可以在 SIMION 中打开这些文件查看电极几何\n")

